/*******************************************************************************
 * MPI自动计时库 - 基于PMPI接口
 * 
 * 原理：MPI标准提供了PMPI_*函数，允许拦截所有MPI_*调用
 * 
 * 编译方法：
 *   mpicxx -shared -fPIC -O2 -o libmpitimer.so libmpitimer.cpp -ldl
 * 
 * 使用方法：
 *   export LD_PRELOAD=./libmpitimer.so
 *   mpirun -np 288 ./yourSolver
 * 
 * 注意：程序需要用 -g 编译以获取更详细的位置信息
 * 
 * 输出：程序结束时自动打印MPI性能报告（含代码位置）
 ******************************************************************************/

#include <mpi.h>
#include <map>
#include <string>
#include <vector>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <execinfo.h>  // backtrace
#include <dlfcn.h>     // dladdr
#include <cxxabi.h>    // abi::__cxa_demangle

// 统计信息结构
struct CallStats {
    long long count;
    double total_time;
    double min_time;
    double max_time;
    long long total_bytes;
    
    CallStats() : count(0), total_time(0), min_time(1e99), max_time(0), total_bytes(0) {}
};

// 代码位置信息
struct LocationInfo {
    std::string function;
    std::string file;
    std::string mpi_call;
    long long count;
    double total_time;
    
    LocationInfo() : count(0), total_time(0) {}
};

// 全局统计数据
static std::map<std::string, CallStats>* g_stats = nullptr;
static std::map<std::string, LocationInfo>* g_location_stats = nullptr;
static int g_rank = -1;
static int g_size = -1;
static bool g_initialized = false;
static double g_program_start = 0;
static bool g_enable_backtrace = false;  // 默认关闭backtrace以提高性能
static int g_sample_rate = 1;  // 采样率：1=全部记录，100=每100次记录1次
static long long g_call_counter = 0;  // 调用计数器，用于采样

// 初始化
static void init_profiler() {
    if (g_stats == nullptr) {
        g_stats = new std::map<std::string, CallStats>();
        g_location_stats = new std::map<std::string, LocationInfo>();
        g_program_start = MPI_Wtime();
        
        // 环境变量：MPITIMER_ENABLE_BACKTRACE=1 开启代码位置追踪（默认关闭）
        const char* enable_bt = getenv("MPITIMER_ENABLE_BACKTRACE");
        if (enable_bt && strcmp(enable_bt, "1") == 0) {
            g_enable_backtrace = true;
        }
        
        // 环境变量：MPITIMER_SAMPLE_RATE=100 表示每100次调用采样1次（默认全部记录）
        const char* sample_env = getenv("MPITIMER_SAMPLE_RATE");
        if (sample_env) {
            int rate = atoi(sample_env);
            if (rate > 0) {
                g_sample_rate = rate;
            }
        }
        
        if (g_rank == -1) {  // 还未初始化时输出提示
            if (g_enable_backtrace) {
                fprintf(stderr, "[MPI Timer] 代码位置追踪已启用（会影响性能）\n");
            }
            if (g_sample_rate > 1) {
                fprintf(stderr, "[MPI Timer] 采样率: 1/%d（减少性能开销）\n", g_sample_rate);
            }
        }
    }
}

// 获取调用者信息（简化版：只获取调用者地址和符号）
static std::string get_caller_info(const char* mpi_name) {
    if (!g_enable_backtrace) return "unknown";
    
    void* buffer[10];
    int nptrs = backtrace(buffer, 10);
    
    if (nptrs < 3) return "unknown";
    
    // buffer[0] = get_caller_info
    // buffer[1] = record_call
    // buffer[2] = MPI_XXX wrapper
    // buffer[3] = actual caller (我们要的)
    
    void* caller_addr = (nptrs > 3) ? buffer[3] : buffer[2];
    
    Dl_info info;
    if (dladdr(caller_addr, &info)) {
        std::string location;
        
        // 尝试demangle函数名
        if (info.dli_sname) {
            int status;
            char* demangled = abi::__cxa_demangle(info.dli_sname, nullptr, nullptr, &status);
            if (status == 0 && demangled) {
                location = demangled;
                free(demangled);
            } else {
                location = info.dli_sname;
            }
        } else {
            location = "unknown_function";
        }
        
        // 添加文件名（如果有）
        if (info.dli_fname) {
            const char* fname = strrchr(info.dli_fname, '/');
            fname = fname ? fname + 1 : info.dli_fname;
            location = std::string(fname) + ":" + location;
        }
        
        return location;
    }
    
    return "unknown";
}

// 记录调用（优化版：支持采样）
static void record_call(const char* name, double time, int bytes = 0) {
    if (!g_initialized) return;
    
    // 记录MPI函数统计（始终记录基本统计）
    CallStats& stats = (*g_stats)[name];
    stats.count++;
    stats.total_time += time;
    if (time < stats.min_time) stats.min_time = time;
    if (time > stats.max_time) stats.max_time = time;
    stats.total_bytes += bytes;
    
    // 记录调用位置统计（按采样率）
    if (g_enable_backtrace) {
        g_call_counter++;
        // 只在采样点才执行昂贵的backtrace操作
        if (g_call_counter % g_sample_rate == 0) {
            std::string caller = get_caller_info(name);
            std::string key = std::string(name) + "@" + caller;
            
            LocationInfo& loc_info = (*g_location_stats)[key];
            loc_info.mpi_call = name;
            loc_info.function = caller;
            loc_info.count++;
            loc_info.total_time += time;
        }
    }
}

// 打印报告
static void print_report() {
    if (!g_initialized || g_rank != 0 || g_stats == nullptr) return;
    
    double program_time = MPI_Wtime() - g_program_start;
    double total_mpi_time = 0;
    
    for (auto& p : *g_stats) {
        total_mpi_time += p.second.total_time;
    }
    
    // 按总时间排序
    std::vector<std::pair<std::string, CallStats>> sorted_stats(g_stats->begin(), g_stats->end());
    std::sort(sorted_stats.begin(), sorted_stats.end(),
              [](const auto& a, const auto& b) { return a.second.total_time > b.second.total_time; });
    
    printf("\n");
    printf("================================================================================\n");
    printf("  MPI Performance Profile (Rank 0, Total Processes: %d)\n", g_size);
    printf("================================================================================\n");
    printf("Program runtime:  %.3f seconds\n", program_time);
    printf("MPI time:         %.3f seconds (%.1f%%)\n", total_mpi_time, total_mpi_time/program_time*100);
    printf("Compute time:     %.3f seconds (%.1f%%)\n", program_time - total_mpi_time, 
           (program_time - total_mpi_time)/program_time*100);
    
    // ========================================================================
    // 表1: 时间统计
    // ========================================================================
    printf("\n");
    printf("[ 表1: MPI调用时间统计 ]\n");
    printf("--------------------------------------------------------------------------------\n");
    printf("%-20s %10s %12s %10s %10s %10s %8s\n", 
           "MPI Call", "Calls", "Total(s)", "Avg(ms)", "Min(ms)", "Max(ms)", "%");
    printf("--------------------------------------------------------------------------------\n");
    
    for (const auto& p : sorted_stats) {
        const std::string& name = p.first;
        const CallStats& stats = p.second;
        double percent = (total_mpi_time > 0) ? (stats.total_time / total_mpi_time * 100) : 0;
        double avg_ms = (stats.count > 0) ? (stats.total_time / stats.count * 1000) : 0;
        
        printf("%-20s %10lld %12.4f %10.3f %10.3f %10.3f %7.1f%%\n",
               name.c_str(), stats.count, stats.total_time, avg_ms,
               stats.min_time * 1000, stats.max_time * 1000, percent);
    }
    printf("--------------------------------------------------------------------------------\n");
    
    // ========================================================================
    // 表2: 数据量统计
    // ========================================================================
    // 筛选出有数据传输的MPI调用
    std::vector<std::pair<std::string, CallStats>> data_stats;
    long long total_bytes = 0;
    for (const auto& p : sorted_stats) {
        if (p.second.total_bytes > 0) {
            data_stats.push_back(p);
            total_bytes += p.second.total_bytes;
        }
    }
    
    if (!data_stats.empty()) {
        printf("\n");
        printf("[ 表2: MPI数据传输统计 ]\n");
        printf("--------------------------------------------------------------------------------\n");
        printf("%-20s %10s %15s %15s %12s\n", 
               "MPI Call", "Calls", "Total Data", "Avg/Call", "Bandwidth");
        printf("--------------------------------------------------------------------------------\n");
        
        for (const auto& p : data_stats) {
            const std::string& name = p.first;
            const CallStats& stats = p.second;
            double total_mb = stats.total_bytes / 1024.0 / 1024.0;
            double avg_kb = (stats.count > 0) ? (stats.total_bytes / stats.count / 1024.0) : 0;
            double bandwidth = (stats.total_time > 0) ? (total_mb / stats.total_time) : 0;
            
            printf("%-20s %10lld %12.2f MB %12.2f KB %9.2f MB/s\n",
                   name.c_str(), stats.count, total_mb, avg_kb, bandwidth);
        }
        
        printf("--------------------------------------------------------------------------------\n");
        printf("Total transferred:   %.2f MB (%.2f GB)\n", 
               total_bytes / 1024.0 / 1024.0, 
               total_bytes / 1024.0 / 1024.0 / 1024.0);
        
        if (total_mpi_time > 0) {
            printf("Average bandwidth:   %.2f MB/s\n", 
                   (total_bytes / 1024.0 / 1024.0) / total_mpi_time);
        }
        printf("--------------------------------------------------------------------------------\n");
    }
    
    // ========================================================================
    // 表3: 热点代码位置统计
    // ========================================================================
    if (g_enable_backtrace && g_location_stats && !g_location_stats->empty()) {
        printf("\n");
        printf("[ 表3: 热点代码位置统计 (Top 20) ]\n");
        printf("--------------------------------------------------------------------------------\n");
        printf("%-20s %10s %12s %8s  %-60s\n", 
               "MPI Call", "Calls", "Total(s)", "%", "调用位置");
        printf("--------------------------------------------------------------------------------\n");
        
        // 按总时间排序
        std::vector<std::pair<std::string, LocationInfo>> sorted_locations(
            g_location_stats->begin(), g_location_stats->end());
        std::sort(sorted_locations.begin(), sorted_locations.end(),
                  [](const auto& a, const auto& b) { 
                      return a.second.total_time > b.second.total_time; 
                  });
        
        int shown = 0;
        for (const auto& p : sorted_locations) {
            if (shown++ >= 20) break;  // 只显示前20个
            
            const LocationInfo& loc = p.second;
            double percent = (total_mpi_time > 0) ? (loc.total_time / total_mpi_time * 100) : 0;
            
            // 智能截断：保留关键部分
            std::string func = loc.function;
            std::string display_func = func;
            
            // 如果太长，只显示最后的函数名（去掉模板参数）
            if (func.length() > 60) {
                // 尝试找到最后一个::后面的内容
                size_t last_colon = func.rfind("::");
                if (last_colon != std::string::npos && last_colon + 2 < func.length()) {
                    std::string short_name = func.substr(last_colon + 2);
                    // 如果还是太长，截断模板参数
                    size_t template_start = short_name.find('<');
                    if (template_start != std::string::npos) {
                        short_name = short_name.substr(0, template_start) + "<...>";
                    }
                    display_func = "..." + short_name;
                    if (display_func.length() > 60) {
                        display_func = display_func.substr(0, 57) + "...";
                    }
                } else {
                    display_func = func.substr(0, 57) + "...";
                }
            }
            
            // 第一行：统计数据
            printf("%-20s %10lld %12.4f %7.1f%%  ",
                   loc.mpi_call.c_str(), loc.count, loc.total_time, percent);
            
            // 完整函数名（可能很长）
            if (func.length() <= 60) {
                printf("%-60s\n", func.c_str());
            } else {
                // 第一行显示简短版本
                printf("%-60s\n", display_func.c_str());
                // 第二行显示完整路径（缩进）
                printf("%54s  └─ %s\n", "", func.c_str());
            }
        }
        
        printf("--------------------------------------------------------------------------------\n");
        printf("说明：此表显示MPI调用最频繁/耗时的代码位置\n");
        printf("      优化建议：重点检查占比>5%%的代码位置\n");
        printf("      完整路径显示在缩进行\n");
        if (g_sample_rate > 1) {
            printf("      注意：使用采样率 1/%d，显示的是采样数据\n", g_sample_rate);
        }
        printf("--------------------------------------------------------------------------------\n");
    }
    
    printf("\n");
    printf("================================================================================\n");
    
    // ========================================================================
    // 诊断建议
    // ========================================================================
    printf("\n");
    printf("[ 诊断建议 ]\n");
    printf("--------------------------------------------------------------------------------\n");
    
    bool has_issue = false;
    
    if (total_mpi_time / program_time > 0.5) {
        printf("  MPI时间占比过高 (%.1f%%)，通信是主要瓶颈\n", 
               total_mpi_time/program_time*100);
        has_issue = true;
    }
    
    for (const auto& p : sorted_stats) {
        if (p.second.total_time / total_mpi_time > 0.3) {
            printf("  %s 是最大瓶颈 (占MPI时间%.1f%%)，建议减少调用频率\n",
                   p.first.c_str(), p.second.total_time / total_mpi_time * 100);
            has_issue = true;
        }
    }
    
    // 检查Allreduce
    auto it = g_stats->find("MPI_Allreduce");
    if (it != g_stats->end() && it->second.count > 1000) {
        printf("  MPI_Allreduce调用 %lld 次，建议增加收敛检查间隔\n", 
               it->second.count);
        has_issue = true;
    }
    
    // 检查Barrier
    it = g_stats->find("MPI_Barrier");
    if (it != g_stats->end() && it->second.count > 100) {
        printf("  MPI_Barrier调用 %lld 次，检查是否有不必要的同步\n", 
               it->second.count);
        has_issue = true;
    }
    
    if (!has_issue) {
        printf("  MPI通信性能良好，无明显瓶颈\n");
    }
    
    printf("--------------------------------------------------------------------------------\n");
    printf("\n");
}

// ============================================================================
// PMPI包装函数 - 拦截常用的MPI调用
// ============================================================================

extern "C" {

// MPI_Init - 初始化profiler
int MPI_Init(int *argc, char ***argv) {
    int ret = PMPI_Init(argc, argv);
    init_profiler();
    PMPI_Comm_rank(MPI_COMM_WORLD, &g_rank);
    PMPI_Comm_size(MPI_COMM_WORLD, &g_size);
    g_initialized = true;
    return ret;
}

// MPI_Init_thread - 初始化profiler（OpenFOAM使用这个）
int MPI_Init_thread(int *argc, char ***argv, int required, int *provided) {
    int ret = PMPI_Init_thread(argc, argv, required, provided);
    init_profiler();
    PMPI_Comm_rank(MPI_COMM_WORLD, &g_rank);
    PMPI_Comm_size(MPI_COMM_WORLD, &g_size);
    g_initialized = true;
    return ret;
}

// MPI_Finalize - 打印报告
int MPI_Finalize() {
    print_report();
    
    if (g_stats) {
        delete g_stats;
        g_stats = nullptr;
    }
    
    if (g_location_stats) {
        delete g_location_stats;
        g_location_stats = nullptr;
    }
    
    return PMPI_Finalize();
}

// MPI_Barrier
int MPI_Barrier(MPI_Comm comm) {
    double start = MPI_Wtime();
    int ret = PMPI_Barrier(comm);
    record_call("MPI_Barrier", MPI_Wtime() - start);
    return ret;
}

// MPI_Bcast
int MPI_Bcast(void *buffer, int count, MPI_Datatype datatype, int root, MPI_Comm comm) {
    int size;
    PMPI_Type_size(datatype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Bcast(buffer, count, datatype, root, comm);
    record_call("MPI_Bcast", MPI_Wtime() - start, count * size);
    return ret;
}

// MPI_Reduce
int MPI_Reduce(const void *sendbuf, void *recvbuf, int count, MPI_Datatype datatype,
               MPI_Op op, int root, MPI_Comm comm) {
    int size;
    PMPI_Type_size(datatype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Reduce(sendbuf, recvbuf, count, datatype, op, root, comm);
    record_call("MPI_Reduce", MPI_Wtime() - start, count * size);
    return ret;
}

// MPI_Allreduce
int MPI_Allreduce(const void *sendbuf, void *recvbuf, int count, MPI_Datatype datatype,
                  MPI_Op op, MPI_Comm comm) {
    int size;
    PMPI_Type_size(datatype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Allreduce(sendbuf, recvbuf, count, datatype, op, comm);
    record_call("MPI_Allreduce", MPI_Wtime() - start, count * size);
    return ret;
}

// MPI_Allgather
int MPI_Allgather(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
                  void *recvbuf, int recvcount, MPI_Datatype recvtype, MPI_Comm comm) {
    int size;
    PMPI_Type_size(sendtype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Allgather(sendbuf, sendcount, sendtype, recvbuf, recvcount, recvtype, comm);
    record_call("MPI_Allgather", MPI_Wtime() - start, sendcount * size);
    return ret;
}

// MPI_Alltoall
int MPI_Alltoall(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
                 void *recvbuf, int recvcount, MPI_Datatype recvtype, MPI_Comm comm) {
    int size;
    PMPI_Type_size(sendtype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Alltoall(sendbuf, sendcount, sendtype, recvbuf, recvcount, recvtype, comm);
    record_call("MPI_Alltoall", MPI_Wtime() - start, sendcount * size);
    return ret;
}

// MPI_Send
int MPI_Send(const void *buf, int count, MPI_Datatype datatype, int dest, int tag, MPI_Comm comm) {
    int size;
    PMPI_Type_size(datatype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Send(buf, count, datatype, dest, tag, comm);
    record_call("MPI_Send", MPI_Wtime() - start, count * size);
    return ret;
}

// MPI_Recv
int MPI_Recv(void *buf, int count, MPI_Datatype datatype, int source, int tag,
             MPI_Comm comm, MPI_Status *status) {
    int size;
    PMPI_Type_size(datatype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Recv(buf, count, datatype, source, tag, comm, status);
    record_call("MPI_Recv", MPI_Wtime() - start, count * size);
    return ret;
}

// MPI_Sendrecv
int MPI_Sendrecv(const void *sendbuf, int sendcount, MPI_Datatype sendtype, int dest, int sendtag,
                 void *recvbuf, int recvcount, MPI_Datatype recvtype, int source, int recvtag,
                 MPI_Comm comm, MPI_Status *status) {
    int size;
    PMPI_Type_size(sendtype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Sendrecv(sendbuf, sendcount, sendtype, dest, sendtag,
                           recvbuf, recvcount, recvtype, source, recvtag, comm, status);
    record_call("MPI_Sendrecv", MPI_Wtime() - start, sendcount * size);
    return ret;
}

// MPI_Wait
int MPI_Wait(MPI_Request *request, MPI_Status *status) {
    double start = MPI_Wtime();
    int ret = PMPI_Wait(request, status);
    record_call("MPI_Wait", MPI_Wtime() - start);
    return ret;
}

// MPI_Waitall
int MPI_Waitall(int count, MPI_Request array_of_requests[], MPI_Status array_of_statuses[]) {
    double start = MPI_Wtime();
    int ret = PMPI_Waitall(count, array_of_requests, array_of_statuses);
    record_call("MPI_Waitall", MPI_Wtime() - start);
    return ret;
}

// MPI_Gather
int MPI_Gather(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
               void *recvbuf, int recvcount, MPI_Datatype recvtype, int root, MPI_Comm comm) {
    int size;
    PMPI_Type_size(sendtype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Gather(sendbuf, sendcount, sendtype, recvbuf, recvcount, recvtype, root, comm);
    record_call("MPI_Gather", MPI_Wtime() - start, sendcount * size);
    return ret;
}

// MPI_Scatter
int MPI_Scatter(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
                void *recvbuf, int recvcount, MPI_Datatype recvtype, int root, MPI_Comm comm) {
    int size;
    PMPI_Type_size(sendtype, &size);
    double start = MPI_Wtime();
    int ret = PMPI_Scatter(sendbuf, sendcount, sendtype, recvbuf, recvcount, recvtype, root, comm);
    record_call("MPI_Scatter", MPI_Wtime() - start, sendcount * size);
    return ret;
}

} // extern "C"

